#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelClipByValue
{
public:
  __aicore__ inline KernelClipByValue() {}
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR clip_value_min, GM_ADDR clip_value_max, GM_ADDR y, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength,
                              uint32_t typeKey)
  {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->typeKey = typeKey;

    


    this->blockLength = blockLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;
    xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x,
                        this->blockLength);
    yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y,
                        this->blockLength);

    minGm.SetGlobalBuffer((__gm__ DTYPE_X *)clip_value_min,
                        this->tileLength);
    this->minValue = *((__gm__ DTYPE_X *)clip_value_min);

    maxGm.SetGlobalBuffer((__gm__ DTYPE_X *)clip_value_max,
                        this->tileLength);
    this->maxValue = *((__gm__ DTYPE_X *)clip_value_max);
    pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(inQueueMIN, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(inQueueMAX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(calcBuf, this->tileLength * sizeof(uint8_t));
    if(typeKey == 3)
      pipe.InitBuffer(tmpBuf, 2 * this->tileLength * sizeof(float));
  }
  
  __aicore__ inline void Process()
  {
    if(typeKey <= 1){
      ProcessMIN();
      ProcessMax();
    }
    
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++)
    {
      if(BUFFER_NUM == 2 && i == loopCount - 2 && lasttileLength <= tileLength){
        continue;
      }
      CopyIn(i);
      Compute(i);
      CopyOut(i);
    }
    if(typeKey <= 1){
      ProcessFree();
    }
  }

private:
  __aicore__ inline void ProcessMIN()
  {
    LocalTensor<DTYPE_X> minLocal = inQueueMIN.AllocTensor<DTYPE_X>();
    Duplicate<DTYPE_X>(minLocal, (DTYPE_X)this->minValue, this->tileLength);
    inQueueMIN.EnQue(minLocal);
  }

  __aicore__ inline void ProcessMax()
  {
    LocalTensor<DTYPE_X> maxLocal = inQueueMAX.AllocTensor<DTYPE_X>();
    Duplicate<DTYPE_X>(maxLocal, (DTYPE_X)this->maxValue, this->tileLength);
    inQueueMAX.EnQue(maxLocal);
  }

  __aicore__ inline void ProcessFree()
  {
    LocalTensor<DTYPE_X> maxLocal = inQueueMAX.DeQue<DTYPE_X>();
    inQueueMAX.FreeTensor(maxLocal);
    maxLocal = inQueueMIN.DeQue<DTYPE_X>();
    inQueueMIN.FreeTensor(maxLocal);
  }

  __aicore__ inline void CopyIn(int32_t progress)
  {
    LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();

    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            xLocal[0],
            xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
      }
      else {
        DataCopy(xLocal[0], xGm[progress * (this->tileLength)],
                 (this->tileLength));
      }
    }
    inQueueX.EnQue(xLocal);
  }
  __aicore__ inline void Compute(int32_t progress)
  {
    if(typeKey == 1){//half
      LocalTensor<half> xLocal = inQueueX.DeQue<half>();
      LocalTensor<half> minLocal = inQueueMIN.DeQue<half>();
      LocalTensor<half> yLocal = outQueueY.AllocTensor<half>();
      LocalTensor<half> maxLocal = inQueueMAX.DeQue<half>();
      LocalTensor<uint8_t> resLocal = calcBuf.Get<uint8_t>();

      Compare(resLocal, xLocal, minLocal, CMPMODE::GT, this->tileLength);
      Select(xLocal, resLocal, xLocal, (half)this->minValue, SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);

      Compare(resLocal, xLocal, maxLocal, CMPMODE::LT, this->tileLength);
      Select(yLocal, resLocal, xLocal, (half)this->maxValue, SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);

      outQueueY.EnQue<half>(yLocal);
      inQueueX.FreeTensor(xLocal);
      inQueueMIN.EnQue<half>(minLocal);
      inQueueMAX.EnQue<half>(maxLocal);
    }else if(typeKey == 0){//float
      LocalTensor<float> xLocal = inQueueX.DeQue<float>();
      LocalTensor<float> minLocal = inQueueMIN.DeQue<float>();
      LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
      LocalTensor<float> maxLocal = inQueueMAX.DeQue<float>();
      LocalTensor<uint8_t> resLocal = calcBuf.Get<uint8_t>();

      Compare(resLocal, xLocal, minLocal, CMPMODE::GT, this->tileLength);
      Select(xLocal, resLocal, xLocal, (float)this->minValue, SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);

      Compare(resLocal, xLocal, maxLocal, CMPMODE::LT, this->tileLength);
      Select(yLocal, resLocal, xLocal, (float)this->maxValue, SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);

      outQueueY.EnQue<float>(yLocal);
      inQueueX.FreeTensor(xLocal);
      inQueueMIN.EnQue<float>(minLocal);
      inQueueMAX.EnQue<float>(maxLocal);
    }else{//int32
      LocalTensor<int32_t> xLocal = inQueueX.DeQue<int32_t>();
      LocalTensor<int32_t> yLocal = outQueueY.AllocTensor<int32_t>();
      LocalTensor<uint8_t> resLocal = calcBuf.Get<uint8_t>();
      LocalTensor<float> tmpLocal = tmpBuf.Get<float>();
      LocalTensor<float> xLocal_f = tmpLocal;
      LocalTensor<float> mLocal_f = tmpLocal[this->tileLength];
      Cast(xLocal_f, xLocal, RoundMode::CAST_FLOOR, this->tileLength);

      Duplicate<float>(mLocal_f, (float)this->minValue, this->tileLength);
      Compare(resLocal, xLocal_f, mLocal_f, CMPMODE::GT, this->tileLength);
      Select(xLocal_f, resLocal, xLocal_f, (float)this->minValue, SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
    
      Duplicate<float>(mLocal_f, (float)this->maxValue, this->tileLength);
      Compare(resLocal, xLocal_f, mLocal_f, CMPMODE::LT, this->tileLength);
      Select(xLocal_f, resLocal, xLocal_f, (float)this->maxValue, SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);

      Cast(yLocal, xLocal_f, RoundMode::CAST_FLOOR, this->tileLength);
      outQueueY.EnQue<int32_t>(yLocal);
      inQueueX.FreeTensor(xLocal);
    }
    
  }
  __aicore__ inline void CopyOut(int32_t progress)
  {
    LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            yLocal, (this->tileLength));
      }

      else {
        DataCopy(yGm[progress * (this->tileLength)], yLocal,
                 (this->tileLength));
      }
    }
    outQueueY.FreeTensor(yLocal);
  }

private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueMIN, inQueueMAX;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
  TBuf<TPosition::VECCALC> calcBuf, tmpBuf;
  GlobalTensor<DTYPE_X> xGm;
  GlobalTensor<DTYPE_X> minGm;
  GlobalTensor<DTYPE_X> maxGm;
  GlobalTensor<DTYPE_Y> yGm;
  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
  uint32_t typeKey;
  uint32_t masklen;
  uint32_t maskCount;
  DTYPE_X minValue;
  DTYPE_X maxValue;
};
extern "C" __global__ __aicore__ void clip_by_value(GM_ADDR x, GM_ADDR clip_value_min, GM_ADDR clip_value_max, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  // TODO: user kernel impl
  KernelClipByValue op;

  op.Init(x, clip_value_min, clip_value_max, y, tiling_data.blockLength,
          tiling_data.tileNum, tiling_data.tileLength,
          tiling_data.lasttileLength, tiling_data.typeKey);
  op.Process();
}
#ifndef __CCE_KT_TEST__
void clip_by_value_do(uint32_t blockDim, void* l2ctrl, void* stream,
                       uint8_t* x, uint8_t* clip_value_min, uint8_t* clip_value_max, uint8_t* y,
                       uint8_t* workspace, uint8_t* tiling) {
  clip_by_value<<<blockDim, l2ctrl, stream>>>(x, clip_value_min, clip_value_max, y, workspace, tiling);
}
#endif
